import socket
import struct
import json

def recv_all(sock, size):
    """可靠接收指定长度的数据"""
    data = b''
    while len(data) < size:
        chunk = sock.recv(min(4096, size - len(data)))
        if not chunk:
            break
        data += chunk
    return data

host = 'localhost'
port = 8080

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect((host, port))
print(f"Connected to server at {host}:{port}")

try:
    while True:
        cmd = input('>> ').strip()
        if not cmd:
            continue
            
        s.sendall(cmd.encode('utf-8'))
        print(f'Sent command: {cmd}')
        
        # 1. 接收报头长度 (4字节大端整数)
        header_size_data = recv_all(s, 4)
        if not header_size_data or len(header_size_data) != 4:
            print("Invalid header size received")
            break
            
        header_size = struct.unpack('>I', header_size_data)[0]
        
        # 2. 接收完整报头
        header_json = recv_all(s, header_size).decode('utf-8')
        if not header_json:
            print("No header received")
            break
            
        try:
            header = json.loads(header_json)
            print("Header received:", header)
        except json.JSONDecodeError:
            print("Invalid JSON header:", header_json)
            break
            
        # 3. 接收完整数据
        total_size = header['total_size']
        raw_data = recv_all(s, total_size)
        
        if len(raw_data) != total_size:
            print(f"Data incomplete: {len(raw_data)}/{total_size} bytes")
            break
            
        # 4. 分离stdout和stderr
        stdout = raw_data[:header['stdout_size']]
        stderr = raw_data[header['stdout_size']:]
        
        # 5. 解码输出
        encoding = header.get('encoding', 'GBK')
        
        print("\n=== STDOUT ===")
        print(stdout.decode(encoding, errors='replace'))
        
        if stderr:
            print("\n=== STDERR ===")
            print(stderr.decode(encoding, errors='replace'))
            
        print(f"\nExit code: {header['exit_code']}")
        
except KeyboardInterrupt:
    print("\nDisconnecting...")
except Exception as e:
    print(f"Error: {str(e)}")
finally:
    s.close()
    print("Connection closed")